'''
breif:
    demo for Cogview-3-flash
    This demo is based on the Cogview-3-flash model and the ZhipuAI API,
    This demo provides user a class to generate and download images based on Cogview-3-flash model.
    For more information about Cogview-3-flash model, please refer to https://bigmodel.cn/dev/api/image-model/cogview
'''
from zhipuai import ZhipuAI
import json,requests
content  =  open('API.json','r',encoding='utf-8').read()
api = json.loads(content)['api']
client = ZhipuAI(api_key=api)

class AI():
    def __init__(self):
        ...
    @staticmethod
    def download(func):
        def wrapper(*args, **kwargs):
            url, download = func(*args, **kwargs)
            if download:
                try:
                    resp = requests.get(url)
                    with open('image.png', 'wb') as f:
                        f.write(resp.content)
                    return "Image downloaded as 'image.png'"
                except Exception as e:
                    return f"Error downloading image: {e}"
            return url
        return wrapper

    ###
    # @prompt: str (discription of the image)
    # @size: str (chose from 1024x1024,768x1344,864x1152,1344x768,1152x864,1440x720,720x1440. default: 1024x1024)
    # @download: bool (download the image or not, default: True)
    # @return: str (url of the image) or download information
    ###
    @download
    def generate(self, prompt, size=None, download=True):
        try:
            response = client.images.generations(
                model="cogview-3-flash",
                prompt=prompt,
                size=size,
            )
            url = response.data[0].url
            return url, download
        except Exception as e:
            raise Exception(f"Error generating image: {e}")
if __name__ == "__main__":
    ai = AI()
    # generate image
    result = ai.generate("a bule cat sleep on the bed",size='1024x1024',download=True)
    print(result)